package online.heycm.excel.common;

import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.stereotype.Component;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.Collection;
import java.util.function.BiFunction;

/**
 * @author hey
 * @version 1.0
 * @since 2024/8/17 20:48
 */
@Component
public class MyBatisBatch {

    public static final int BATCH_SIZE = 1000;
    private static SqlSessionFactory sqlSessionFactory;

    public MyBatisBatch(SqlSessionFactory sqlSessionFactory) {
        MyBatisBatch.sqlSessionFactory = sqlSessionFactory;
    }

    public static <T, U, R> int execute(Collection<T> data, Class<U> mapperClass, BiFunction<T, U, R> function) {
        int i = 1;
        SqlSession session = null;
        try {
            session = sqlSessionFactory.openSession(ExecutorType.BATCH);
            U mapper = session.getMapper(mapperClass);
            int size = data.size();
            for (T element : data) {
                function.apply(element, mapper);
                if (i % BATCH_SIZE == 0 || i == size) {
                    session.flushStatements();
                }
                i++;
            }
            session.commit(!TransactionSynchronizationManager.isSynchronizationActive());
        } catch (Exception e) {
            e.printStackTrace();
            if (session != null) {
                session.rollback();
            }
        } finally {
            if (session != null) {
                session.close();
            }
        }
        return i - 1;
    }
}
